# 导入所需库
import websocket  # 用于WebSocket通信
import uuid       # 生成唯一客户端ID
import json       # 处理JSON数据
import urllib.request  # 发送HTTP请求
import urllib.parse    # 构造URL参数
import requests
from config.host import server_address, client_id
import os
import io # 新增导入

def queue_prompt(prompt):
    """提交提示到服务器并返回执行ID"""
    # 构造包含提示和客户端ID的请求体
    p = {"prompt": prompt, "client_id": client_id}
    # 发送POST请求到/prompt端点
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request(
        f"http://{server_address}/prompt", 
        data=data
    )
    return json.loads(urllib.request.urlopen(req).read())
 
def get_image(filename, subfolder, folder_type):
    """根据参数获取生成的图像二进制数据"""
    # 构造URL查询参数
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    # 发送GET请求获取图像
    with urllib.request.urlopen(f"http://{server_address}/view?{url_values}") as response:
        return response.read()
    
def upload_image_with_formdata(url, file_path, form_fields=None):
    """
    使用urllib实现formdata格式文件上传
    :param url: 目标接口地址
    :param file_path: 本地图片路径
    :param form_fields: 其他表单字段字典
    """
    boundary = '----' + str(uuid.uuid4())
    
    # 构建form-data数据体
    data = []
    
    # 添加普通表单字段
    if form_fields:
        for key, value in form_fields.items():
            data.append(f'--{boundary}\r\n'.encode())
            data.append(f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode())
            data.append(f'{value}\r\n'.encode())
    
    # 添加文件字段
    with open(file_path, 'rb') as f:
        file_content = f.read()
    filename = file_path.split('/')[-1]
    data.append(f'--{boundary}\r\n'.encode())
    data.append(f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'.encode())
    data.append(b'Content-Type: image/jpeg\r\n\r\n')
    data.append(file_content)
    data.append(f'\r\n--{boundary}--\r\n'.encode())
    
    # 合并数据
    body = b''.join(data)
    
    # 构建请求
    headers = {
        'Content-Type': f'multipart/form-data; boundary={boundary}',
        'Content-Length': str(len(body))
    }
    
    request = urllib.request.Request(url, data=body, headers=headers, method='POST')
    
    try:
        with urllib.request.urlopen(request) as response:
            return response.read().decode()
    except urllib.error.URLError as e:
        return f"Upload failed: {e.reason}"


def upload_image(file_or_buffer, filename=None):
    """
    上传图片到 ComfyUI 服务器。
    支持文件路径 (str) 或内存中的 buffer (io.BytesIO)。
    
    :param file_or_buffer: 图片的文件路径或 BytesIO buffer。
    :param filename: 当上传 buffer 时，必须提供此文件名。
    :return: 包含状态码和响应内容的字典。
    """
    url = f"http://{server_address}/upload/image"
    
    # 判断输入是文件路径还是 buffer
    if isinstance(file_or_buffer, str):
        # 输入是文件路径
        if not filename:
            filename = os.path.basename(file_or_buffer)
        with open(file_or_buffer, 'rb') as f:
            files = {'image': (filename, f)}
            # 增加 overwrite=true 可以在同名时覆盖，避免因临时文件重名导致错误
            data = {'type': 'input', 'overwrite': 'true'}
            response = requests.post(url, files=files, data=data)
    elif isinstance(file_or_buffer, io.BytesIO):
        # 输入是 buffer
        if not filename:
            raise ValueError("上传 buffer 时必须提供 'filename' 参数。")
        file_or_buffer.seek(0) # 确保从头读取
        files = {'image': (filename, file_or_buffer)}
        data = {'type': 'input', 'overwrite': 'true'}
        response = requests.post(url, files=files, data=data)
    else:
        raise TypeError("不支持的 'file_or_buffer' 类型。请提供 str 或 io.BytesIO。")

    return {
        'status': response.status_code,
        'content': response.text
    }

def get_history(prompt_id):
    """获取指定提示ID的执行历史记录"""
    with urllib.request.urlopen(f"http://{server_address}/history/{prompt_id}") as response:
        return json.loads(response.read())
 
def clean_dir():
    """删除服务器上的临时文件"""
    with urllib.request.urlopen(f"http://{server_address}/api/clean_usedata") as response:
        return json.loads(response.read())

def get_images(ws, prompt):
    """通过WebSocket监听执行过程并获取最终生成的图像"""
    # 提交提示并获取执行ID
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}
    
    # 持续监听WebSocket消息
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            # 检测执行完成信号
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break  # 执行完成时退出循环
        else:
            # 二进制数据流（预览图）处理逻辑（已跳过）
            continue
    
    # 获取执行历史并提取生成的图像
    history = get_history(prompt_id)[prompt_id]
    for node_id in history['outputs']:
        node_output = history['outputs'][node_id]
        images_output = []
        if 'images' in node_output:
            for image in node_output['images']:
                # 获取每个图像的二进制数据
                image_data = get_image(image['filename'], image['subfolder'], image['type'])
                images_output.append(image_data)
        output_images[node_id] = images_output
    
    return output_images
 
# 定义工作流提示模板




